# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import unittest

from coin.exchange.base.ws.subscriber import get_next_unit_time


class SubscriberUtilTest(unittest.TestCase):
  def test_get_next_unit_time(self):
    current_dt = datetime.datetime(2017, 12, 28, 13, 5, 27)
    self.assertEqual(datetime.datetime(2017, 12, 28, 14),
                     get_next_unit_time(datetime.timedelta(hours=1), current_dt=current_dt))
    self.assertEqual(datetime.datetime(2017, 12, 28, 13, 6),
                     get_next_unit_time(datetime.timedelta(minutes=1), current_dt=current_dt))
    self.assertEqual(
        datetime.datetime(2017, 12, 28, 13, 30),
        get_next_unit_time(datetime.timedelta(hours=1),
                           datetime.timedelta(minutes=30),
                           current_dt=current_dt))

    current_dt = datetime.datetime(2017, 12, 28, 23, 59, 59)
    self.assertEqual(datetime.datetime(2017, 12, 29),
                     get_next_unit_time(datetime.timedelta(hours=1), current_dt=current_dt))
    self.assertEqual(datetime.datetime(2017, 12, 29),
                     get_next_unit_time(datetime.timedelta(minutes=1), current_dt=current_dt))
    self.assertEqual(
        datetime.datetime(2017, 12, 29, 0, 30),
        get_next_unit_time(datetime.timedelta(hours=1),
                           datetime.timedelta(minutes=30),
                           current_dt=current_dt))


if __name__ == "__main__":
  unittest.main()
